#include "sampler.hh"

int sampler_t::operator()(const sampler_t::vec_t& prob)
{
  return sample(prob);
}

int sampler_t::sample(const sampler_t::vec_t& prob)
{
#ifdef DEBUG
  assert_msg(prob.size() > 0, "|prob| > 0");
  for(size_t j=0; j<prob.size(); ++j)
  assert_msg(prob[j]>=0., "non-negative number");
#endif
  size_t K = prob.size();
  double tot = std::accumulate(prob.begin(), prob.end(), 0.);
  double u = unif(RandGen) * tot;

  double cum = 0.;
  int ret = 0;
  for (size_t j = 0; j < K; ++j)
  {
    cum += prob.at(j);
    if (u <= cum)
    {
      ret = j;
      break;
    }
  }
  return ret;
}

int sampler_t::log_sample(const sampler_t::vec_t& lprob)
{
#ifdef DEBUG
  assert_msg(lprob.size() > 0, "|lprob| > 0");
#endif
  double maxval = *std::max_element(lprob.begin(), lprob.end());

  size_t K = lprob.size();
  double tot = 0.;
  for (size_t j = 0; j < K; ++j)
  {
    tot += std::exp(lprob.at(j) - maxval);
  }

  double u = unif(RandGen) * tot;
  double cum = 0.;
  int ret = 0;
  for (size_t j = 0; j < K; ++j)
  {
    cum += std::exp(lprob.at(j) - maxval);
    if (u <= cum)
    {
      ret = j;
      break;
    }
  }
  return ret;
}

